import numpy as np
import pandas as pd
import seaborn as sns


def get_columns_name_with_na(df):
    sum = df.isna().sum(axis=0)
    idx = sum > 0
    return df.columns[idx], sum[idx]


if '__main__' == __name__:
    pd.set_option('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', 1000, 'display.expand_frame_repr', False)

    df = sns.load_dataset('titanic')
    print(df[:5])

    cols, na_nums = get_columns_name_with_na(df)
    print(cols)
    print(na_nums)

    from sklearn.impute import SimpleImputer
    df['age'] = SimpleImputer(strategy='mean').fit_transform(df['age'].values.reshape(-1, 1))

    cols, na_nums = get_columns_name_with_na(df)
    print(cols)
    print(na_nums)

    df['embarked'] = SimpleImputer(strategy='most_frequent').fit_transform(df['embarked'].values.reshape(-1, 1))

    cols, na_nums = get_columns_name_with_na(df)
    print(cols)
    print(na_nums)
